#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rospy
import numpy as np
from sensor_msgs.msg import PointCloud2, PointField
import sensor_msgs.point_cloud2 as pc2
from std_msgs.msg import Header
from visualization_msgs.msg import MarkerArray, Marker
from geometry_msgs.msg import Point
# import pcl
# import pcl.pcl_visualization
from collections import defaultdict

class GridRefinementNode:
    def __init__(self):
        rospy.init_node('grid_refinement_node', anonymous=True)
        
        # 订阅patchwork的地面点和非地面点
        self.ground_sub = rospy.Subscriber('/ground_segmentation/ground', PointCloud2, self.ground_callback)
        self.nonground_sub = rospy.Subscriber('/segmented_cloud_pure1', PointCloud2, self.nonground_callback)
        
        # 发布重新分类后的点云
        self.refined_ground_pub = rospy.Publisher('/refined_ground', PointCloud2, queue_size=1)
        self.refined_nonground_pub = rospy.Publisher('/refined_nonground', PointCloud2, queue_size=1)
        
        # 发布网格高度可视化 - 使用MarkerArray
        self.grid_height_pub = rospy.Publisher('/grid_height_visualization', MarkerArray, queue_size=1)
        
        # 网格参数
        self.grid_resolution = 1.0  # 1米网格
        self.grid_bounds = {
            'x_min': -50.0, 'x_max': 50.0,
            'y_min': -50.0, 'y_max': 50.0
        }
        
        # 存储数据
        self.ground_points = None
        self.nonground_points = None
        self.ground_grid = defaultdict(list)  # 存储每个网格内的地面点
        self.grid_avg_heights = {}  # 存储每个网格的平均高度
        
        rospy.loginfo("Grid Refinement Node initialized")
        
    def ground_callback(self, msg):
        """处理地面点云数据"""
        start_time = rospy.Time.now()
        
        # 计算点的数量
        point_count = sum(1 for _ in pc2.read_points(msg, skip_nans=True))
        rospy.loginfo("Received ground points: %d", point_count)
        
        self.ground_points = self.pointcloud2_to_array(msg)
        self.process_ground_points()
        
        end_time = rospy.Time.now()
        processing_time = (end_time - start_time).to_sec()
        rospy.loginfo("Ground points processing time: %.3f seconds", processing_time)
        
    def nonground_callback(self, msg):
        """处理非地面点云数据"""
        start_time = rospy.Time.now()
        
        # 计算点的数量
        point_count = sum(1 for _ in pc2.read_points(msg, skip_nans=True))
        rospy.loginfo("Received nonground points: %d", point_count)
        
        self.nonground_points = self.pointcloud2_to_array(msg)
        self.process_nonground_points()
        
        end_time = rospy.Time.now()
        processing_time = (end_time - start_time).to_sec()
        rospy.loginfo("Nonground points processing time: %.3f seconds", processing_time)
        
    def pointcloud2_to_array(self, msg):
        """将PointCloud2消息转换为numpy数组"""
        points = []
        for point in pc2.read_points(msg, field_names=("x", "y", "z", "intensity"), skip_nans=True):
            points.append([point[0], point[1], point[2], point[3]])
        return np.array(points)
        
    def get_grid_index(self, x, y):
        """获取点所在的网格索引"""
        grid_x = int(np.floor(x / self.grid_resolution))
        grid_y = int(np.floor(y / self.grid_resolution))
        return (grid_x, grid_y)
        
    def process_ground_points(self):
        """处理地面点，构建网格"""
        if self.ground_points is None or len(self.ground_points) == 0:
            return
            
        # 清空之前的网格数据
        self.ground_grid.clear()
        self.grid_avg_heights.clear()
        
        # 将地面点分配到网格
        grid_start_time = rospy.Time.now()
        for point in self.ground_points:
            x, y, z = point[0], point[1], point[2]
            grid_idx = self.get_grid_index(x, y)
            self.ground_grid[grid_idx].append(z)
        
        grid_end_time = rospy.Time.now()
        grid_time = (grid_end_time - grid_start_time).to_sec()
        rospy.loginfo("Grid assignment time: %.3f seconds", grid_time)
            
        # 计算每个网格的平均高度
        height_start_time = rospy.Time.now()
        for grid_idx, heights in self.ground_grid.items():
            if len(heights) > 0:
                self.grid_avg_heights[grid_idx] = np.mean(heights)
        
        height_end_time = rospy.Time.now()
        height_time = (height_end_time - height_start_time).to_sec()
        rospy.loginfo("Height calculation time: %.3f seconds", height_time)
                
        rospy.loginfo("Processed ground points into %d grids", len(self.grid_avg_heights))
        
    def process_nonground_points(self):
        """处理非地面点，基于网格高度进行重新分类"""
        if self.nonground_points is None or len(self.nonground_points) == 0:
            return
            
        if not self.grid_avg_heights:
            rospy.logwarn("No ground grid data available, skipping nonground processing")
            return
            
        refined_ground = []
        refined_nonground = []
        
        # 按网格分组非地面点
        grouping_start_time = rospy.Time.now()
        grid_nonground_points = {}
        for point in self.nonground_points:
            x, y, z = point[0], point[1], point[2]
            grid_idx = self.get_grid_index(x, y)
            
            if grid_idx not in grid_nonground_points:
                grid_nonground_points[grid_idx] = []
            grid_nonground_points[grid_idx].append(z)
        
        grouping_end_time = rospy.Time.now()
        grouping_time = (grouping_end_time - grouping_start_time).to_sec()
        rospy.loginfo("Grid grouping time: %.3f seconds", grouping_time)
        
        # 重新分类非地面点
        classification_start_time = rospy.Time.now()
        for point in self.nonground_points:
            x, y, z = point[0], point[1], point[2]
            grid_idx = self.get_grid_index(x, y)
            
            # 检查该网格是否有地面高度数据
            if grid_idx in self.grid_avg_heights:
                grid_height = self.grid_avg_heights[grid_idx] + 0.1
                
                # 如果点的高度低于网格平均高度+0.2，重新分类为地面点
                if z <= grid_height:
                    refined_ground.append([x, y, z, 50.0])  # 强度设为50
                else:
                    refined_nonground.append([x, y, z, 100.0])  # 强度设为100
            else:
                # 如果网格没有地面数据，使用该网格内非地面点的最低点+0.2作为过滤阈值
                if grid_idx in grid_nonground_points and len(grid_nonground_points[grid_idx]) > 0:
                    grid_min_height = min(grid_nonground_points[grid_idx])
                    grid_threshold = grid_min_height + 0.1
                    
                    if z <= grid_threshold:
                        refined_ground.append([x, y, z, 50.0])  # 强度设为50
                    else:
                        refined_nonground.append([x, y, z, 100.0])  # 强度设为100
                else:
                    # 如果该网格内没有非地面点，保持为非地面点
                    refined_nonground.append([x, y, z, 100.0])
        
        classification_end_time = rospy.Time.now()
        classification_time = (classification_end_time - classification_start_time).to_sec()
        rospy.loginfo("Point classification time: %.3f seconds", classification_time)
                
        # 将原始地面点也加入，强度设为50
        if self.ground_points is not None:
            for point in self.ground_points:
                refined_ground.append([point[0], point[1], point[2], 50.0])
                
        # 发布重新分类后的点云
        self.publish_pointcloud(refined_ground, self.refined_ground_pub, "base_laser", "refined_ground")
        self.publish_pointcloud(refined_nonground, self.refined_nonground_pub, "base_laser", "refined_nonground")
        
        # 发布网格高度可视化
        self.publish_grid_height_visualization()
        
        rospy.loginfo("Refined classification: %d ground points, %d nonground points", 
                     len(refined_ground), len(refined_nonground))
        
    def publish_pointcloud(self, points, publisher, frame_id, topic_name):
        """发布点云消息"""
        if len(points) == 0:
            return
            
        # 创建PointCloud2消息
        header = Header()
        header.stamp = rospy.Time.now()
        header.frame_id = frame_id
        
        # 定义点云字段
        fields = [
            PointField('x', 0, PointField.FLOAT32, 1),
            PointField('y', 4, PointField.FLOAT32, 1),
            PointField('z', 8, PointField.FLOAT32, 1),
            PointField('intensity', 12, PointField.FLOAT32, 1)
        ]
        
        # 创建点云消息
        cloud_msg = pc2.create_cloud(header, fields, points)
        publisher.publish(cloud_msg)
        
    def publish_grid_height_visualization(self):
        """发布网格高度可视化 - 只显示有地面数据的网格，空网格不可视化"""
        if not self.grid_avg_heights:
            return
            
        # 创建MarkerArray
        marker_array = MarkerArray()
        marker_id = 0
        
        # 获取所有网格的坐标范围
        grid_coords = list(self.grid_avg_heights.keys())
        if not grid_coords:
            return
            
        min_x = min(coord[0] for coord in grid_coords)
        max_x = max(coord[0] for coord in grid_coords)
        min_y = min(coord[1] for coord in grid_coords)
        max_y = max(coord[1] for coord in grid_coords)
        
        # 只遍历有地面数据的网格
        for (grid_x, grid_y), height in self.grid_avg_heights.items():
                
                # 计算网格边界坐标
                x_min = grid_x * self.grid_resolution
                x_max = (grid_x + 1) * self.grid_resolution
                y_min = grid_y * self.grid_resolution
                y_max = (grid_y + 1) * self.grid_resolution
                
                # 创建网格边缘Marker（空心网格）
                self.create_hollow_grid_marker(marker_array, marker_id, x_min, x_max, y_min, y_max, height)
                marker_id += 1
                
        # 发布MarkerArray
        self.grid_height_pub.publish(marker_array)
        

            
    def create_hollow_grid_marker(self, marker_array, marker_id, x_min, x_max, y_min, y_max, height):
        """创建空心网格Marker（只画边缘）"""
        # 创建4条边的Marker
        edges = [
            # 下边
            {'pos': [(x_min + x_max)/2, y_min, height/2], 'scale': [x_max - x_min, 0.05, max(0.1, height + 0.5)]},
            # 上边
            {'pos': [(x_min + x_max)/2, y_max, height/2], 'scale': [x_max - x_min, 0.05, max(0.1, height + 0.5)]},
            # 左边
            {'pos': [x_min, (y_min + y_max)/2, height/2], 'scale': [0.05, y_max - y_min, max(0.1, height + 0.5)]},
            # 右边
            {'pos': [x_max, (y_min + y_max)/2, height/2], 'scale': [0.05, y_max - y_min, max(0.1, height + 0.5)]}
        ]
        
        for i, edge in enumerate(edges):
            marker = Marker()
            marker.header.frame_id = "base_laser"
            marker.header.stamp = rospy.Time.now()
            marker.ns = "grid_height"
            marker.id = marker_id * 4 + i  # 每个网格4条边，ID不重复
            marker.type = Marker.CUBE
            marker.action = Marker.ADD
            
            # 设置边缘位置和尺寸
            marker.pose.position.x = edge['pos'][0]
            marker.pose.position.y = edge['pos'][1]
            marker.pose.position.z = edge['pos'][2]
            marker.pose.orientation.w = 1.0
            
            # 设置边缘尺寸
            marker.scale.x = edge['scale'][0]
            marker.scale.y = edge['scale'][1]
            marker.scale.z = edge['scale'][2]
            
            # 根据高度设置颜色（高度越高，颜色越红）
            height_normalized = max(0.0, min(1.0, (height + 2.0) / 4.0))
            marker.color.r = height_normalized
            marker.color.g = 1.0 - height_normalized
            marker.color.b = 0.0
            marker.color.a = 0.8  # 稍微提高透明度
            
            marker_array.markers.append(marker)
        
    def run(self):
        """运行节点"""
        rospy.loginfo("Grid Refinement Node is running...")
        rospy.spin()

if __name__ == '__main__':
    try:
        node = GridRefinementNode()
        node.run()
    except rospy.ROSInterruptException:
        pass 